package com.shujia.spark.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel, NaiveBayes, NaiveBayesModel}
import org.apache.spark.ml.linalg
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo02PersonIndexModelTrain {
  def main(args: Array[String]): Unit = {
    /**
     * 机器学习的一般流程
     */
    val spark: SparkSession = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName.replace("$", ""))
      .master("local")
      .config("spark.sql.shuffle.partitions", "2")
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._

    val personDF: DataFrame = spark
      .read
      .format("libsvm")
      .load("spark/data/mllib/data/personIndex.txt")

    // 1、加载数据，进行数据特征工程处理，将数据变成MLLib模型所能够接收的向量

    // 2、切分数据集：80%作为训练集（训练模型）、20%作为测试集（评估模型）
    val splitArr: Array[Dataset[Row]] = personDF.randomSplit(Array(0.8, 0.2))
    val trainDF: Dataset[Row] = splitArr(0)
    val testDF: Dataset[Row] = splitArr(1)

    /**
     * 3、选择合适的模型：
     * 数据集有无Label：
     * 有->有监督学习->Label是否连续：
     * 连续的：回归算法模型
     * 离散的：分类算法模型
     * 无->无监督学习
     *
     * 人体指标数据中，有Label，并且是离散的 -> 逻辑回归模型进行分类
     */
    val lr: LogisticRegression = new LogisticRegression()
      .setMaxIter(10) // 最大的迭代次数
      .setRegParam(0.3) // 正则化参数
      .setElasticNetParam(0.8)

    // Fit the model，将训练集带入模型进行训练
    val lrModel: LogisticRegressionModel = lr.fit(trainDF)

    //    // Train a NaiveBayes model.
    //    val model: NaiveBayesModel = new NaiveBayes()
    //      //      .setFeaturesCol("features")
    //      //      .setLabelCol("label")
    //      .fit(trainDF)
    //
    //    // Select example rows to display.
    //    val predictions: DataFrame = model.transform(testDF)
    //    predictions.show()
    //    println(s"贝叶斯模型的准确率为：${1 - predictions.where($"label" =!= $"prediction").count() / testDF.count().toDouble}")


    // 4、用测试集评估模型
    val testTransformDF: DataFrame = lrModel.transform(testDF)
    println(s"逻辑回归模型的准确率为：${1 - testTransformDF.where($"label" =!= $"prediction").count() / testDF.count().toDouble}")

    // 5、如果模型通过评估，则可以保存模型进行应用
    lrModel
      .write
      .overwrite()
      .save("spark/data/mllib/model/person")


    // 1、加载新的数据集，进行同样的数据特征工程处理，将每条数据转换成向量的格式
    val personRDD: RDD[String] = spark.sparkContext.textFile("spark/data/mllib/data/test/person.txt")
    val newPersonDF: DataFrame = personRDD
      .map(line => {
        val arr: Array[Double] = line.split(" ").map(_.toDouble)
        Tuple1(Vectors.dense(arr))
      })
      // 将RDD转换成DF，并将列命名为features
      .toDF("features")

    // 2、加载训练好的模型
    //    val lrModel: LogisticRegressionModel = LogisticRegressionModel
    //      .read
    //      .load("spark/data/mllib/model/person")

    // 3、将处理好的数据带入模型

    // 对一批数据进行预测
    val predictionDF: DataFrame = lrModel.transform(newPersonDF)
    predictionDF.show()

    // 将数据转换成向量，通过predict方法即可对一条数据进行预测
    val personVec: linalg.Vector = Vectors.dense(Array(3.9, 3.3, 2.2, 105.6, 69.6, 56.0, 87))
    val res: Double = lrModel.predict(personVec)
    println(res)
  }

}
